from identification import identify_mrgwl_event, mrgwl_case_identification
from dataLoader import VwndLoader
import xarray as xr
from mrgLike import MrgwLikeEvents, mrgw_like_list_to_nc, mrgw_like_list_to_json
from pathlib import Path
import pandas as pd
import numpy as np

def identify():
    """加载数据并识别 MRGWL 事件"""

    # 请根据你的数据路径与格式调整数据读取方式
    # 这里使用了一个在 dataLoader 中定义的 VwndLoader 类来按照月份加载数据
    for vwnd_file in VwndLoader():

        # 按照月份读取滤波好的经向风数据
        vwnd = xr.open_dataarray(vwnd_file)

        # 识别 MRGWL 事件
        mrg = identify_mrgwl_event(vwnd)

        # 缓存数据（非必须）
        file_name = vwnd_file.name
        mrg.save(f"cache/{file_name}.feather")
        print(f"Saved {file_name} to cache.")


def connect():
    """连接 MRGWL 事件"""

    # 加载缓存
    folder = Path("cache")
    mrgl_events = MrgwLikeEvents()
    files = list(folder.glob("*.feather"))
    files.sort()
    for file in files:
        mrgl_events.load_and_concat(file)
    
    # 连接 MRGWL 事件
    mrgwl = mrgwl_case_identification(mrgl_events, 0.2)

    # 保存结果
    mrgw_like_list_to_nc(mrgwl, "dist/mrgwl_1979-2022_0.2.nc")
    mrgw_like_list_to_json(mrgwl, "dist/mrgwl_1979-2022_0.2.json")

    # 连接 MRGWL 事件
    mrgwl = mrgwl_case_identification(mrgl_events, 0.06)

    # 保存结果
    mrgw_like_list_to_nc(mrgwl, "dist/mrgwl_1979-2022_0.06.nc")
    mrgw_like_list_to_json(mrgwl, "dist/mrgwl_1979-2022_0.06.json")


def check():
    """画图检查识别结果"""
    from matplotlib import pyplot as plt
    data = xr.open_dataset("dist/mrgwl_1979-2022_0.1.nc")
    file = VwndLoader().path / "FLTv.850.200008.nc"
    vwnd = xr.open_dataarray(file)
    time_range = (pd.Timestamp("2000-08-01"), pd.Timestamp("2000-08-31"))

    # 选出时间
    t0 = data.time.values[:, 0]
    idx = np.where((t0 >= time_range[0]) & (t0 <= time_range[1]))[0]
    mrgwl = data.isel(id=idx)

    fig = plt.figure(figsize=(6, 9))
    fig.subplots_adjust(left=0.17, right=0.96, top=0.99, bottom=0.05)
    ax = fig.add_subplot(111)
    ax.set_ylim(*time_range)
    ax.set_xlim(0, 360)
    ax.set_xlabel("Longitude")

    ax.contourf(vwnd.longitude, vwnd.time, 
                vwnd.sel(latitude=slice(5, -5)).mean("latitude"), 
                levels=np.linspace(-5, 5, 21), cmap="RdBu_r", 
                extend='both')
    for i in range(len(mrgwl.id)):
        ax.plot(mrgwl.longitude[i], mrgwl.time[i], "k.-", markersize=2, linewidth=0.5)
        ax.text(mrgwl.longitude[i][0], mrgwl.time[i][0], mrgwl.id.values[i], fontsize=6, ha="left", va="top")

    plt.savefig("doc/chech.png", dpi=150)

if __name__ == "__main__":
    # identify()
    connect()
    # check()
